# Nuclear Engine Maya Skeleton Export
# Code by Rockie X.Lee

# Usage:
#   import maya.cmds as cmds
#   cmds.loadPlugin('NEPATH/NEMayaSkeletonExport.py')
#   cmds.neMayaSkeletonExport()

import sys
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx
import maya.OpenMayaAnim as OpenMayaAnim

class NEKeyFrame():
    def __init__(self):
        self.__name = ""
        self.__time = 0.0
        self.__position = OpenMaya.MVector()
        self.__oritation = OpenMaya.MQuaternion()
        self.__scale = OpenMaya.MVector(1.0, 1.0, 1.0)

    def getName(self):
        return self.__name

    def getTime(self):
        return self.__time;

    def getPosition(self):
        return self.__position

    def getOritation(self):
        return self.__oritation

    def getScale(self):
        return self.__scale

    def setData(self, time, dagPath):
        joint = OpenMayaAnim.MFnIkJoint(dagPath)
        self.__name = joint.name()
        self.__time = time
        print str(self.__name)
        self.__position = joint.getTranslation(OpenMaya.MSpace.kPostTransform)
        # print self.__position.x, self.__position.y, self.__position.z
        joint.getRotation(self.__oritation, OpenMaya.MSpace.kPostTransform)
        # print self.__oritation.x, self.__oritation.y, self.__oritation.z, self.__oritation.w
        """
        scalePtr = OpenMaya.MScriptUtil().asDoublePtr()
        joint.getScale(scalePtr)
        self.__scale.x = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 0)
        self.__scale.y = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 1)
        self.__scale.z = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 2)
        """
        # print self.__scale.x, self.__scale.y, self.__scale.z


class NEJoint():
    def __init__(self):
        self.__dagPath = 0
        self.__handle = 0xFFFF
        self.__parent = 0xFFFF
        self.__name = ""
        self.__parentName = ""
        self.__position = OpenMaya.MVector()
        self.__oritation = OpenMaya.MQuaternion()
        self.__scale = OpenMaya.MVector(1.0, 1.0, 1.0)

    def getDagPath(self):
        return self.__dagPath

    def getHandle(self):
        return self.__handle

    def setParent(self, parent):
        self.__parent = parent

    def getParent(self):
        return self.__parent

    def getName(self):
        return self.__name

    def getParentName(self):
        return self.__parentName

    def getPosition(self):
        return self.__position

    def getOritation(self):
        return self.__oritation

    def getScale(self):
        return self.__scale

    def setData(self, handle, dagPath):
        self.__dagPath = dagPath
        joint = OpenMayaAnim.MFnIkJoint(dagPath)
        self.__handle = handle
        self.__name = joint.name()
        print str(self.__name)
        if joint.parentCount() > 0:
            parent = joint.parent(0)
            if parent.hasFn(OpenMaya.MFn.kJoint):
                parentJoint = OpenMayaAnim.MFnIkJoint(parent)
                self.__parentName = parentJoint.name()
                # print self.__parentName
        self.__position = joint.getTranslation(OpenMaya.MSpace.kPostTransform)
        # print self.__position.x, self.__position.y, self.__position.z
        joint.getRotation(self.__oritation, OpenMaya.MSpace.kPostTransform)
        # print self.__oritation.x, self.__oritation.y, self.__oritation.z, self.__oritation.w
        """
        scalePtr = OpenMaya.MScriptUtil().asDoublePtr()
        joint.getScale(scalePtr)
        self.__scale.x = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 0)
        self.__scale.y = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 1)
        self.__scale.z = OpenMaya.MScriptUtil.getDoubleArrayItem(scalePtr, 2)
        """
        # print self.__scale.x, self.__scale.y, self.__scale.z


neMayaSkeletonPluginCmdName = "neMayaSkeletonExport"

# Command
class NEMayaSkeletonExportCommand(OpenMayaMPx.MPxCommand):
    def __init__(self):
        OpenMayaMPx.MPxCommand.__init__(self)
        self.__dir = ""
        self.__name = "Skeleton"
        self.__numJoints = 0
        self.__joints = []
        self.__jointHandleMap = dict()
        self.__animation = dict()
        self.__animationLength = 1.0

    def doIt(self, argList):
        print "[Begin] Nuclear Engine Maya Skeleton Export===================================="
        self.exportSkeleton()
        if self.__numJoints > 0:
            self.exportAnimation()
        print "[Success] Nuclear Engine Maya Skeleton Export=================================="

    def exportSkeleton(self):
        # Create a selection list and iterator
        selectList = OpenMaya.MSelectionList()
        OpenMaya.MGlobal.getActiveSelectionList(selectList)
        iter = OpenMaya.MItSelectionList(selectList)
        dagIter = OpenMaya.MItDag(OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)
        while not iter.isDone():
            objectPath = OpenMaya.MDagPath()
            iter.getDagPath(objectPath)
            dagIter.reset(objectPath.node(), OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)
            while not dagIter.isDone():
                dagPath = OpenMaya.MDagPath()
                dagIter.getPath(dagPath)
                if dagPath.hasFn(OpenMaya.MFn.kJoint):
                    joint = NEJoint()
                    joint.setData(self.__numJoints, dagPath)
                    self.__joints.append(joint)
                    self.__jointHandleMap[joint.getName()] = self.__numJoints
                    self.__numJoints += 1
                dagIter.next()
            iter.next()

        if self.__numJoints > 0:
            self.setJointsParentHandle()
            self.writeSkeletonData(self.__dir)
        else:
            print "No seleted joints !"

    def exportAnimation(self):
        curTime = 0
        while curTime < 24.0 * self.__animationLength:
            curTime += 1
            OpenMayaAnim.MAnimControl.setCurrentTime(OpenMaya.MTime(curTime, OpenMaya.MTime.kFilm));
            for joint in self.__joints:
                keyFrame = NEKeyFrame()
                keyFrame.setData(curTime / 24.0, joint.getDagPath())
                self.__animation[keyFrame.getName()].append(keyFrame)
        self.writeAnimationData(self.__dir)

    def setJointsParentHandle(self):
        for joint in self.__joints:
            self.__animation[joint.getName()] = []
            parentName = joint.getParentName()
            if parentName != "":
                if self.__jointHandleMap.has_key(parentName):
                    joint.setParent(self.__jointHandleMap[parentName])

    def writeSkeletonData(self, dir):
        file = open(dir + str(self.__name) + ".skeletonpy", "w");
        file.write("Name = " + str(self.__name) + "\n")
        file.write("NumJoints = " + str(self.__numJoints) + "\n")
        for joint in self.__joints:
            file.write("Handle = " + str(joint.getHandle()) + "\n")
            file.write("Parent = " + str(joint.getParent()) + "\n")
            file.write("Name = " + joint.getName() + "\n")
            position = joint.getPosition()
            file.write("Position = " + str(position.x) + " " + str(position.y) + " " + str(position.z) + "\n")
            oritation = joint.getOritation()
            file.write("Oritation = " + str(oritation.x) + " " + str(oritation.y) + " " + str(oritation.z) + " " + str(oritation.w) + "\n")
            scale = joint.getScale()
            file.write("Scale = " + str(scale.x) + " " + str(scale.y) + " " + str(scale.z) + "\n")
        file.close()
        print "Finished export skeleton data !"

    def writeAnimationData(self, dir):
        file = open(dir + str(self.__name) + ".animationpy", "w");
        file.write("Name = " + str(self.__name) + "\n")
        file.write("Length = " + str(self.__animationLength) + "\n")
        file.write("NumNodeTracks = " + str(len(self.__animation)) + "\n")
        for key, track in self.__animation.iteritems():
            file.write("Handle = " + str(self.__jointHandleMap[key]) + "\n")
            file.write("NumKeyFrames = " + str(len(track)) + "\n")
            for keyFrame in track:
                file.write("Time = " + str(keyFrame.getTime()) + "\n")
                position = keyFrame.getPosition()
                file.write("Position = " + str(position.x) + " " + str(position.y) + " " + str(position.z) + "\n")
                oritation = keyFrame.getOritation()
                file.write("Oritation = " + str(oritation.x) + " " + str(oritation.y) + " " + str(oritation.z) + " " + str(oritation.w) + "\n")
                scale = keyFrame.getScale()
                file.write("Scale = " + str(scale.x) + " " + str(scale.y) + " " + str(scale.z) + "\n")
        file.close()
        print "Finished export animation data !"

# Creator
def cmdCreator():
    return OpenMayaMPx.asMPxPtr(NEMayaSkeletonExportCommand())


# Syntax creator
def syntaxCreator():
    syntax = OpenMaya.MSyntax()
    return syntax


# Initialize the plug-in
def initializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.registerCommand(neMayaSkeletonPluginCmdName, cmdCreator, syntaxCreator)
    except:
        sys.stderr.write("Failed to register command: %s\n" % neMayaSkeletonPluginCmdName)
        raise


# Uninitialize the plug-in
def uninitializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.deregisterCommand(neMayaSkeletonPluginCmdName)
    except:
        sys.stderr.write("Failed to unregister command: %s\n" % neMayaSkeletonPluginCmdName)
        raise

